import redis
import json
from datetime import datetime, timedelta
import math

# --- 配置 ---
REDIS_HOST = 'localhost'
REDIS_PORT = 6379
# 模拟的起始和结束日期，必须与 precompute_cache.py 中的数据范围一致
SIMULATION_START_DATE = datetime(2021, 5, 1)
SIMULATION_END_DATE = datetime(2021, 5, 15) # 范围是 [start, end)，所以结束日期要加一天
NUM_COMMUNITIES = 74
# 我们将把计算结果保存到这个 JSON 文件中
OUTPUT_FILE = "initial_bike_counts.json"

def calculate_and_save_initial_counts():
    """
    连接 Redis，计算每个社区的理论最小初始单车数，并保存到 JSON 文件。
    """
    print("开始连接 Redis...")
    try:
        r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0, decode_responses=True)
        r.ping()
        print("✅ Redis 连接成功！")
        
    except Exception as e:
        print(f"❌ Redis 连接失败: {e}")
        return

    initial_counts = {}
    print("\n开始计算每个社区的初始单车数，这可能需要一点时间...")

    # --- 外层循环：遍历每一个社区 ---
    for community_id in range(1, NUM_COMMUNITIES + 1):
        print(f"  - 正在处理社区 {community_id}/{NUM_COMMUNITIES}...")
        
        cumulative_net_flow = 0
        max_deficit = 0  # 记录最大的“赤字”（即累积净流量的最低点）
        
        redis_key = f"community:{community_id}:netflow"
        
        # --- 内层循环：按时间顺序遍历所有时间槽 ---
        current_slot = SIMULATION_START_DATE
        while current_slot < SIMULATION_END_DATE:
            # 将时间格式化为与 Redis 中存储的 key 一致的字符串
            slot_str = current_slot.isoformat()
            
            # 从 Redis 获取这个时间槽的净流量
            net_flow_str = r.hget(redis_key, slot_str)
            
            # 如果 Redis 中没有这个槽的数据，说明净流量为 0
            net_flow = int(net_flow_str) if net_flow_str is not None else 0
            
            # 累加净流量
            cumulative_net_flow += net_flow
            
            # 更新最大赤字
            if cumulative_net_flow < max_deficit:
                max_deficit = cumulative_net_flow
            
            # 前进到下一个30分钟时间槽
            current_slot += timedelta(minutes=30)
            
        # 理论最小初始单车数 = -max_deficit
        # 例如，如果最大赤字是 -58，那么初始值就需要是 58，才能保证在那个点车辆数不小于0
        theoretical_minimum = -max_deficit
        initial_counts[community_id] = theoretical_minimum

    print("\n✅ 所有社区计算完成！")
    
    # --- 保存结果到 JSON 文件 ---
    output_path = f"../data/processed/{OUTPUT_FILE}"
    with open(output_path, 'w') as f:
        json.dump(initial_counts, f, indent=2)
        
    print(f"结果已成功保存到: {output_path}")
    print("\n部分结果预览:")
    for i in range(1, min(6, NUM_COMMUNITIES + 1)):
        print(f"  - 社区 {i}: {initial_counts.get(str(i), 'N/A')} 辆")

if __name__ == "__main__":
    calculate_and_save_initial_counts()